package queue

import (
	"context"
	"time"

	"github.com/go-kratos/kratos/v2/log"
)

// RejectHandler Reject push into queue if return false.
type RejectHandler func(ctx context.Context) bool

type Queue interface {
	Push(value interface{})
	TryPush(value interface{}, timeout time.Duration) bool
	Poll() interface{}
	TryPoll(timeout time.Duration) (interface{}, bool)
	size() int
}

type QueueOptions struct {
	Limit int
}

type blockingQueue struct {
	q       chan interface{} // 阻塞队列使用通道表示
	limit   int              // 阻塞队列的大小
	ctx     context.Context
	handler RejectHandler
	log     *log.Helper
}

func NewBlockingQueue(options *QueueOptions, logger log.Logger) (q Queue) {
	q = newBlockingQueue(logger, context.Background(), options)
	return
}

func newBlockingQueue(logger log.Logger, ctx context.Context, options *QueueOptions) (q *blockingQueue) {
	q = &blockingQueue{
		q:     make(chan interface{}, options.Limit),
		limit: options.Limit,
		ctx:   ctx,
		log:   log.NewHelper(logger),
	}
	q.SetRejectHandler(func(ctx context.Context) bool {
		q.log.Error("Pool Reject")
		return false
	})
	return
}

// SetRejectHandler Set a reject handler.
func (q *blockingQueue) SetRejectHandler(handler RejectHandler) {
	q.handler = handler
}

func (q *blockingQueue) Push(value interface{}) {
	ok := true
	if q.handler != nil {
		select {
		case q.q <- value:
			return
		default:
			ok = q.handler(q.ctx)
		}
	}
	if ok {
		q.q <- value
	}
}

func (q *blockingQueue) TryPush(value interface{}, timeout time.Duration) bool {
	ctx, cancel := context.WithTimeout(q.ctx, timeout)
	defer cancel()
	select {
	case q.q <- value:
		return true
	case <-ctx.Done():
	}
	return false
}

func (q *blockingQueue) Poll() interface{} {
	ret := <-q.q
	return ret
}

func (q *blockingQueue) TryPoll(timeout time.Duration) (ret interface{}, ok bool) {
	ctx, cancel := context.WithTimeout(q.ctx, timeout)
	defer cancel()
	select {
	case ret = <-q.q:
		return ret, true
	case <-ctx.Done():
	}
	return ret, false
}

func (q *blockingQueue) size() int {
	return len(q.q)
}
